#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import random
import torch
import math 
import copy
from typing import Union, List
 
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
 
 
def scale_K_list(K_list, target_sum):
    import math
    K_sum = sum(K_list)
    if K_sum <= target_sum:
        return K_list
    scale = target_sum / K_sum
    K_temp = [max(1, math.floor(k * scale)) for k in K_list]
    diff = target_sum - sum(K_temp)
    # 按最大误差分配
    errors = [(k * scale - math.floor(k * scale), idx) for idx, k in enumerate(K_list)]
    if diff > 0:
        errors.sort(reverse=True)
        for i in range(diff):
            K_temp[errors[i][1]] += 1
    elif diff < 0:
        errors.sort()
        i = 0
        while diff < 0 and i < len(K_list):
            if K_temp[errors[i][1]] > 1:
                K_temp[errors[i][1]] -= 1
                diff += 1
            i += 1
    return K_temp
 
 
@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_V4_single_mulit")
class AscendGroupedMatmulV4SingleMulit(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
 
        return input_case
 
    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1]
        biasOptional = case_config.inputs[2]
        KTling = case_config.inputs[16]
        # print('======x-weight================================')
        # print(KTling)
        # print("weight[0].shape", weight[0].shape)
        N = weight[0].shape[1]
        x_dtype = x[0].dtype
        #biasOptional的长度和weight相同，内部的tensor的shape是N
        biasOptional = [copy.deepcopy(biasOptional[0])  for _ in range(len(weight))]
        KTling = [copy.deepcopy(KTling[0])  for _ in range(len(weight))]
        #x为FLOAT16、weight为FLOAT16、biasOptional为FLOAT16   x为BFLOAT16、weight为BFLOAT16、biasOptional为FLOAT32
        K_list= [] 
        for i in range(len(weight)):
            weight[i].dtype = x_dtype
            weight[i].shape[-1] = N 
            K_list.append(weight[i].shape[0])
            biasOptional[i].shape[0] = N
            if x_dtype == 'bf16':
                biasOptional[i].dtype = 'fp32'
 
        # K_sum = sum(K_list)
        # if K_sum >= 13107:  
        #     scale = 13107/ K_sum
        #     K_list = [math.floor(k * scale) for k in K_list]
        #     for i in range(len(K_list)):
        #         weight[i].shape[0] = K_list[i]
        #     K_sum = sum(K_list)
        K_list = scale_K_list(K_list, 13107)
        for i in range(len(K_list)):
            weight[i].shape[0] = K_list[i]
            if weight[i].shape[0] ==1 and weight[i].shape[0] == 1:
                weight[i].shape[0] += 1
                K_list[i] += 1 
        K_sum = sum(K_list)
        extra_K = random.choice([0,10,20,30])
        K = K_sum + extra_K
        x[0].shape[0] = K
        for i in range(len(KTling)):
            KTling[i].range_values = K_list[i]
 
        case_config.inputs[2] = biasOptional
        case_config.inputs[16] = KTling
        if case_config.inputs[0][0].dtype == "bf16":
            case_config.inputs[17].range_values = "bf16"
        elif case_config.inputs[0][0].dtype == "fp16":
            case_config.inputs[17].range_values = "fp16"
        else:
            case_config.inputs[17].range_values = "fp32"
        # print('======x-weight================================')
        # print(f'kkkkkkkkkkkkkkkkkkkkK{K}{K_list}')
        # print(x)
        # print(weight)
        # print(case_config.inputs[2])
        # print(case_config.inputs[-1])
 
        return case_config